package com.geor.gcv.ict.classifier;

import com.geor.gcv.core.mat.Mats;
import com.geor.gcv.core.process.GrayProcess;
import org.nutz.json.Json;
import org.opencv.core.Mat;
import org.tensorflow.*;
import org.tensorflow.types.UInt8;

import java.util.Arrays;
import java.util.List;

public class TensorFlows {

    static float[][] scale = new float[64][64];

    static {
        for (int i = 0; i < scale.length; i++) {
            for (int j = 0; j < scale[i].length; j++) {
                scale[i][j] = 255f;
            }
        }
    }

    // public static Tensor<Float> doToTensorData(List<Mat> mats) {
    //// byte[][] imageBytes = Mats.MatToByte(mat, ".png");
    // try (Graph g = new Graph()) {
    // GraphBuilder b = new GraphBuilder(g);
    // final int H = mat.height();
    // final int W = mat.width();
    // final float scale = 255;
    //
    // final Output<String> input = b.constant("images_bacth", imageBytes);
    // final Output<Float> output = b.div(b.resizeBilinear(
    // b.expandDims(b.cast(b.decodeJpeg(input, 1), Float.class),
    // b.constant("make_batch", 0)),
    // b.constant("size", new int[] { H, W })), b.constant("scale", scale));
    // try (Session s = new Session(g)) {
    // return
    // s.runner().fetch(output.op().name()).run().get(0).expect(Float.class);
    // }
    // }
    // }

    public static Tensor<Float> doToTensorData(Mat mat) {
        return doToTensorDatas(Arrays.asList(mat));
    }

    public static Tensor<Float> doToTensorDatas(List<Mat> mats) {
        byte[][] imageBytesBacth = new byte[mats.size()][0];
        for (int i = 0; i < imageBytesBacth.length; i++) {
            imageBytesBacth[i] = Mats.MatToByte(mats.get(i), ".png");
        }
        final int H = mats.get(0).height();
        final int W = mats.get(0).width();
        final float scale = 255;
        float[][][][] matirx = new float[imageBytesBacth.length][H][W][1];

        for (int i = 0; i < imageBytesBacth.length; i++) {
            try (Graph g = new Graph()) {
                GraphBuilder b = new GraphBuilder(g);
                try (Session s = new Session(g)) {
                    byte[] imageBytes = imageBytesBacth[i];
                    final Output<String> input = b.constant("images_bacth", imageBytes);
                    final Output<Float> output = b.div(
                            b.resizeBilinear(b.expandDims(b.cast(b.decodeJpeg(input, 1), Float.class),
                                    b.constant("make_batch", 0)), b.constant("size", new int[]{H, W})),
                            b.constant("scale", scale));
                    Tensor<Float> tf = s.runner().fetch(output.op().name()).run().get(0).expect(Float.class);
                    float[][][][] matirx1 = new float[1][H][W][1];
                    tf.copyTo(matirx1);
                    for (int n = 0; n < matirx1[0].length; n++) {
                        for (int m = 0; m < matirx1[0][0].length; m++) {
                            matirx[i][n][m][0] = matirx1[0][n][m][0];
                        }
                    }
                }
            }
        }

        Tensor<Float> tf = Tensors.create(matirx);
        return tf;
    }

    public static Tensor<?> bulidTensor(byte[] matirx, int width, int height) {
        float[] matirxf = new float[matirx.length];
        for (int i = 0; i < matirx.length; i++) {
            matirxf[i] = matirx[i];
        }
        return bulidTensor(matirxf, width, height);
    }

    public static Tensor<?> bulidTensor(float[] matirx, int width, int height) {
        float maxtirxs[][] = new float[width][height];
        for (int i = 0; i < matirx.length; i++) {
            maxtirxs[i / width][i % width] = matirx[i];
        }
        return bulidTensor(maxtirxs);
    }

    public static Tensor<?> bulidTensor(float[][] matirx) {
        try (Graph g = new Graph(); Session s = new Session(g)) {
            Output<?> x = g.opBuilder("Placeholder", "images_bacth").setAttr("shape", Shape.make(-1, 64))
                    .setAttr("dtype", DataType.FLOAT).setAttr("make_batch", 0).setAttr("channels", 1).build().output(0);// 创建一个OP
            Output<?> x1 = g.opBuilder("Placeholder", "ts").setAttr("dtype", DataType.FLOAT).build().output(0);// 创建一个OP
            g.opBuilder("Div", "scale").addInput(x).addInput(x1).build().output(0);

            // 图片 除号

            try (Tensor<Float> tx = Tensor.create(matirx, Float.class);
                 Tensor<Float> ts = Tensor.create(scale, Float.class);
                 Tensor<?> output = s.runner().feed("images_bacth", tx).feed("ts", ts).fetch("scale").run().get(0)) {
                // float[][] matirx1 = new float[64][64];
                // output.copyTo(matirx1);
                // for (int i = 0; i < matirx1.length; i++) {
                // for (int j = 0; j < matirx1[i].length; j++) {
                // System.out.print(matirx1[i][j] + " ");
                // }
                // System.out.println();
                // }
                // System.out.println(Json.toJson(output.shape()));
                return output;
            }
        }
    }

    public static void printTensors(Mat mat) {
        byte[] imageBytes = Mats.MatToByte(mat, ".png");

        @SuppressWarnings("unchecked")
        Tensor<Float> image = (Tensor<Float>) bulidTensor(imageBytes, mat.width(), mat.height());
        long[] rshape = image.shape();
        System.out.println(Json.toJson(rshape));
        float[][][][] matirx = new float[1][(int) rshape[1]][(int) rshape[2]][1];
        image.copyTo(matirx);
        System.out.println(Json.toJson(matirx));
    }

    static class GraphBuilder {
        GraphBuilder(Graph g) {
            this.g = g;
        }

        Output<Float> div(Output<Float> x, Output<Float> y) {
            return binaryOp("Div", x, y);
        }

        <T> Output<T> sub(Output<T> x, Output<T> y) {
            return binaryOp("Sub", x, y);
        }

        <T> Output<Float> resizeBilinear(Output<T> images, Output<Integer> size) {
            return binaryOp3("ResizeBilinear", images, size);
        }

        <T> Output<T> expandDims(Output<T> input, Output<Integer> dim) {
            return binaryOp3("ExpandDims", input, dim);
        }

        <T, U> Output<U> cast(Output<T> value, Class<U> type) {
            DataType dtype = DataType.fromClass(type);
            return g.opBuilder("Cast", "Cast").addInput(value).setAttr("DstT", dtype).build().<U>output(0);
        }

        Output<UInt8> decodeJpeg(Output<String> contents, long channels) {
            return g.opBuilder("DecodeJpeg", "DecodeJpeg").addInput(contents).setAttr("channels", channels).build()
                    .<UInt8>output(0);
        }

        <T> Output<T> constant(String name, Object value, Class<T> type) {
            try (Tensor<T> t = Tensor.<T>create(value, type)) {
                return g.opBuilder("Const", name).setAttr("dtype", DataType.fromClass(type)).setAttr("value", t).build()
                        .<T>output(0);
            }
        }

        Output<String> constant(String name, byte[] value) {
            return this.constant(name, value, String.class);
        }

        Output<Integer> constant(String name, int value) {
            return this.constant(name, value, Integer.class);
        }

        Output<Integer> constant(String name, int[] value) {
            return this.constant(name, value, Integer.class);
        }

        Output<Float> constant(String name, float value) {
            return this.constant(name, value, Float.class);
        }

        private <T> Output<T> binaryOp(String type, Output<T> in1, Output<T> in2) {
            return g.opBuilder(type, type).addInput(in1).addInput(in2).build().<T>output(0);
        }

        private <T, U, V> Output<T> binaryOp3(String type, Output<U> in1, Output<V> in2) {
            return g.opBuilder(type, type).addInput(in1).addInput(in2).build().<T>output(0);
        }

        private Graph g;
    }

    public static Tensor<Float> doToTensorData1(Mat mat) {
        mat = new GrayProcess(mat).process();
        byte[] imageBytes = Mats.MatToByte(mat, ".jpg");

        try (Graph g = new Graph()) {
            GraphBuilder b = new GraphBuilder(g);
            final int H = 64;
            final int W = 64;
            final float mean = 1;
            final float scale = 255f;

            final Output<String> input = b.constant("images_bacth", imageBytes);
            final Output<Float> output = b
                    .div(b.sub(
                            b.resizeBilinear(b.expandDims(b.cast(b.decodeJpeg(input, 1), Float.class),
                                    b.constant("make_batch", 0)), b.constant("size", new int[]{H, W})),
                            b.constant("mean", mean)), b.constant("scale", scale));
            try (Session s = new Session(g)) {
                return s.runner().fetch(output.op().name()).run().get(0).expect(Float.class);
            }
        }
    }

}
